{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jeffheaton/app_deep_learning/blob/main/t81_558_class_04_4_batch_norm.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# T81-558: Applications of Deep Neural Networks\n",
    "\n",
    "**Module 4: Training for Tabular Data**\n",
    "\n",
    "- Instructor: [Jeff Heaton](https://sites.wustl.edu/jeffheaton/), McKelvey School of Engineering, [Washington University in St. Louis](https://engineering.wustl.edu/Programs/Pages/default.aspx)\n",
    "- For more information visit the [class website](https://sites.wustl.edu/jeffheaton/t81-558/).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 4 Material\n",
    "\n",
    "- Part 4.1: Using K-Fold Cross-validation with PyTorch [[Video]](https://www.youtube.com/watch?v=Q8ZQNvZwsNE&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/app_deep_learning/blob/main/t81_558_class_04_1_kfold.ipynb)\n",
    "- Part 4.2: Training Schedules for PyTorch  [[Video]](https://www.youtube.com/watch?v=lMMlbmfvKDQ&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/app_deep_learning/blob/main/t81_558_class_04_2_schedule.ipynb)\n",
    "- Part 4.3: Dropout Regularization [[Video]](https://www.youtube.com/watch?v=4ixjgw6Q42U&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/app_deep_learning/blob/main/t81_558_class_04_3_dropout.ipynb)\n",
    "- **Part 4.4: Batch Normalization** [[Video]](https://www.youtube.com/watch?v=1U5nOKh9OLQ&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/app_deep_learning/blob/main/t81_558_class_04_4_batch_norm.ipynb)\n",
    "- Part 4.5: RAPIDS for Tabular Data [[Video]](https://www.youtube.com/watch?v=KgoXuhG_kfs&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/app_deep_learning/blob/main/t81_558_class_04_5_rapids.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Google CoLab Instructions\n",
    "\n",
    "The following code ensures that Google CoLab is running and maps Google Drive if needed. We also initialize the PyTorch device to either GPU/MPS (if available) or CPU.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: not using Google CoLab\n",
      "Using device: mps\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "import torch\n",
    "\n",
    "try:\n",
    "    import google.colab\n",
    "\n",
    "    COLAB = True\n",
    "    print(\"Note: using Google CoLab\")\n",
    "except:\n",
    "    print(\"Note: not using Google CoLab\")\n",
    "    COLAB = False\n",
    "\n",
    "# Make use of a GPU or MPS (Apple) if one is available.  (see module 3.2)\n",
    "device = (\n",
    "    \"mps\"\n",
    "    if getattr(torch, \"has_mps\", False)\n",
    "    else \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"cpu\"\n",
    ")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "\n",
    "# Early stopping (see module 3.4)\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=5, min_delta=0, restore_best_weights=True):\n",
    "        self.patience = patience\n",
    "        self.min_delta = min_delta\n",
    "        self.restore_best_weights = restore_best_weights\n",
    "        self.best_model = None\n",
    "        self.best_loss = None\n",
    "        self.counter = 0\n",
    "        self.status = \"\"\n",
    "\n",
    "    def __call__(self, model, val_loss):\n",
    "        if self.best_loss is None:\n",
    "            self.best_loss = val_loss\n",
    "            self.best_model = copy.deepcopy(model.state_dict())\n",
    "        elif self.best_loss - val_loss >= self.min_delta:\n",
    "            self.best_model = copy.deepcopy(model.state_dict())\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "            self.status = f\"Improvement found, counter reset to {self.counter}\"\n",
    "        else:\n",
    "            self.counter += 1\n",
    "            self.status = f\"No improvement in the last {self.counter} epochs\"\n",
    "            if self.counter >= self.patience:\n",
    "                self.status = f\"Early stopping triggered after {self.counter} epochs.\"\n",
    "                if self.restore_best_weights:\n",
    "                    model.load_state_dict(self.best_model)\n",
    "                return True\n",
    "        return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4.4: Batch Normalization\n",
    "\n",
    "In this section, we're going to take a deep dive into an advanced concept in training neural networks known as Batch Normalization. The goal of batch normalization is to accelerate learning and stabilize the learning process. the same dataset used previously in this chapter. To implement batch normalization we will modify both preprocessing and a neural network setup. The original preprocessing method used a z-score standardization to make the data more suited for the training process. The model itself was a simple feed-forward neural network built with PyTorch's nn.Sequential API.\n",
    "\n",
    "In the revised code, we've made two major changes. The first change was in data preprocessing. Since batch normalization can reduce the impact of input distribution changes (referred to as internal covariate shift), we can remove the step of z-score normalization. Batch normalization tends to make the network less sensitive to the scale and distribution of its inputs, thereby minimizing the need for manual, meticulous data normalization.\n",
    "\n",
    "The second, and most crucial, change was made in the architecture of the neural network itself. We've inserted batch normalization layers into our model by using the nn.BatchNorm1d() function. It's important to note that the batch normalization layers are typically added after the linear (or convolutional for ConvNets) layers but before the activation function. In our case, the sequence is: Linear -> BatchNorm -> ReLU.\n",
    "\n",
    "The batch normalization layers normalize the activations and gradients propagating through a neural network, making the model training more efficient. This can even have a slight regularization effect, somewhat akin to Dropout.\n",
    "\n",
    "Remember that the use of batch normalization may require some additional computational resources due to the additional complexity of the model, but it often results in a significant performance boost that more than compensates for the extra computation time.\n",
    "\n",
    "In summary, the introduction of batch normalization in our neural network model simplifies preprocessing and can potentially improve model training speed, stability, and overall performance. Now, let's dig into the code and see how these enhancements are implemented in practice. \n",
    "\n",
    "The modified code follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fold #1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jeff/miniconda3/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py:1950: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  if not trace_joint and original_inpt.storage().size() != updated_inpt.storage().size():\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 27/500, Validation Loss: 184.96066284179688, Early stopping triggered after 5 epochs.\n",
      "Fold #2\n",
      "Epoch 23/500, Validation Loss: 301.9339294433594, Early stopping triggered after 5 epochs.\n",
      "Fold #3\n",
      "Epoch 27/500, Validation Loss: 935.7130126953125, Early stopping triggered after 5 epochs.\n",
      "Fold #4\n",
      "Epoch 29/500, Validation Loss: 446.3973388671875, Early stopping triggered after 5 epochs.\n",
      "Fold #5\n",
      "Epoch 31/500, Validation Loss: 937.2984008789062, Early stopping triggered after 5 epochs.\n",
      "Fold score (RMSE): 4.957376956939697\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from sklearn.model_selection import KFold\n",
    "\n",
    "# Read the data set\n",
    "df = pd.read_csv(\n",
    "    \"https://data.heatonresearch.com/data/t81-558/jh-simple-dataset.csv\",\n",
    "    na_values=['NA','?'])\n",
    "\n",
    "# Generate dummies for job\n",
    "df = pd.concat([df,pd.get_dummies(df['job'],prefix=\"job\",dtype=int)],axis=1)\n",
    "df.drop('job', axis=1, inplace=True)\n",
    "\n",
    "# Generate dummies for area\n",
    "df = pd.concat([df,pd.get_dummies(df['area'],prefix=\"area\",dtype=int)],axis=1)\n",
    "df.drop('area', axis=1, inplace=True)\n",
    "\n",
    "# Generate dummies for product\n",
    "df = pd.concat([df,pd.get_dummies(df['product'],prefix=\"product\",dtype=int)],axis=1)\n",
    "df.drop('product', axis=1, inplace=True)\n",
    "\n",
    "# Missing values for income\n",
    "med = df['income'].median()\n",
    "df['income'] = df['income'].fillna(med)\n",
    "\n",
    "# Convert to PyTorch Tensors\n",
    "x_columns = df.columns.drop(['age', 'id'])\n",
    "x = torch.tensor(df[x_columns].values, dtype=torch.float32, device=device)\n",
    "y = torch.tensor(df['age'].values, dtype=torch.float32, device=device).view(-1, 1)\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# Cross-Validate\n",
    "kf = KFold(n_splits=5, shuffle=True, random_state=42)\n",
    "\n",
    "# Early stopping parameters\n",
    "patience = 10\n",
    "\n",
    "fold = 0\n",
    "for train_idx, test_idx in kf.split(x):\n",
    "    fold += 1\n",
    "    print(f\"Fold #{fold}\")\n",
    "\n",
    "    x_train, x_test = x[train_idx], x[test_idx]\n",
    "    y_train, y_test = y[train_idx], y[test_idx]\n",
    "\n",
    "    # PyTorch DataLoader\n",
    "    train_dataset = TensorDataset(x_train, y_train)\n",
    "    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "\n",
    "    # Create the model and optimizer\n",
    "    model = nn.Sequential(\n",
    "        nn.Linear(x.shape[1], 20),\n",
    "        nn.BatchNorm1d(20),  # BatchNorm layer\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(20, 10),\n",
    "        nn.BatchNorm1d(10),  # BatchNorm layer\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(10, 1)\n",
    "    )\n",
    "    model = torch.compile(model,backend=\"aot_eager\").to(device)\n",
    "\n",
    "    optimizer = optim.Adam(model.parameters())\n",
    "    loss_fn = nn.MSELoss()\n",
    "\n",
    "    # Early Stopping variables\n",
    "    best_loss = float('inf')\n",
    "    early_stopping_counter = 0\n",
    "\n",
    "    # Training loop\n",
    "    EPOCHS = 500\n",
    "    epoch = 0\n",
    "    done = False\n",
    "    es = EarlyStopping()\n",
    "\n",
    "    while not done and epoch<EPOCHS:\n",
    "        epoch += 1\n",
    "        model.train()\n",
    "        for x_batch, y_batch in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            output = model(x_batch)\n",
    "            loss = loss_fn(output, y_batch)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        # Validation\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            val_output = model(x_test)\n",
    "            val_loss = loss_fn(val_output, y_test)\n",
    "\n",
    "        if es(model, val_loss):\n",
    "            done = True\n",
    "\n",
    "    print(f\"Epoch {epoch}/{EPOCHS}, Validation Loss: \"\n",
    "      f\"{val_loss.item()}, {es.status}\")\n",
    "\n",
    "# Final evaluation\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    oos_pred = model(x_test)\n",
    "score = torch.sqrt(loss_fn(oos_pred, y_test)).item()\n",
    "print(f\"Fold score (RMSE): {score}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Batch normalization can often improve the performance of a model, but it's not a guarantee. It can depend on many factors such as the specific problem being addressed, the data, the architecture of the model, and the specific training regime.\n",
    "\n",
    "* Dataset: Batch normalization is particularly useful when dealing with high dimensional data and tends to be more effective with larger datasets. If your dataset is small or simple, batch normalization may not make a significant difference and might even cause the performance to decrease slightly.\n",
    "\n",
    "* Model Complexity: In simpler models, batch normalization may not lead to a significant performance boost, and might even degrade the performance slightly, as it introduces extra complexity and computation.\n",
    "\n",
    "* Training Parameters: The improvement due to batch normalization also depends on the other hyperparameters you're using, like the learning rate and the batch size. If the original model was already well-optimized with these parameters, the addition of batch normalization may not help much, and could potentially even harm performance.\n",
    "\n",
    "In other words, while batch normalization is generally a useful technique for improving model performance and stability, it's not a silver bullet and it won't always lead to an improvement. It's always a good idea to experiment with different architectures and techniques to see what works best for your specific problem."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3.9 (torch)",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
